package chainx

import (
	"fmt"
	"strings"

	"github.com/tmc/langchaingo/outputparser"
	"github.com/tmc/langchaingo/prompts"
)

const (
	_destinations = "destinations"
	_input        = "input"
	_formatting   = "formatting"

	MULTI_PROMPT_ROUTER_TEMPLATE = `Given a raw text input to a language model select the model prompt best suited for
the input. You will be given the names of the available prompts and a description of
what the prompt is best suited for. You may also revise the original input if you
think that revising it will ultimately lead to a better response from the language
model.

<< FORMATTING >>
Return a markdown code snippet with a JSON object formatted to look like:
{{.formatting}}

REMEMBER: "destination" MUST be one of the candidate prompt names specified below OR \
it can be "DEFAULT" if the input is not well suited for any of the candidate prompts.
REMEMBER: "next_inputs" can just be the original input if you don't think any \
modifications are needed.

<< CANDIDATE PROMPTS >>
{{.destinations}}

<< INPUT >>
{{.input}}`
)

var _outputParser = outputparser.NewStructured([]outputparser.ResponseSchema{
	{
		Name:        _destinations,
		Description: "name of the prompt to use or \"DEFAULT\"",
	},
	{
		Name:        _input,
		Description: "对用户输入的原始提示词进行优化：",
	},
})

func createPrompt(handlers []Handler) prompts.PromptTemplate {
	return prompts.PromptTemplate{
		Template:       MULTI_PROMPT_ROUTER_TEMPLATE,
		InputVariables: []string{_input},
		TemplateFormat: prompts.TemplateFormatGoTemplate,
		PartialVariables: map[string]any{
			"destinations": handlerDestinations(handlers),
			"model_names":  handleNames(handlers),
			"formatting":   _outputParser.GetFormatInstructions(),
		},
	}
}

func handleNames(handlers []Handler) string {
	var hn strings.Builder
	for i, handle := range handlers {
		if i > 0 {
			hn.WriteString(", ")
		}
		hn.WriteString(handle.Name())
	}

	return hn.String()
}

func handlerDestinations(handlers []Handler) string {
	var hs strings.Builder
	for _, handle := range handlers {
		hs.WriteString(fmt.Sprintf("- %s: %s\n", handle.Name(), handle.Description()))
	}

	return hs.String()
}
